import math

import matplotlib as mpl
from matplotlib import pyplot as plt
from sklearn.metrics import r2_score

import cdm
import fine_tune
import getData
import tradaboost_r2

mpl.rcParams['font.sans-serif'] = ['SimHei']
mpl.rcParams['axes.unicode_minus'] = False


def error(y_predict, y_test):  # 定义计算误差平方和函数，，，传入的是估算出的值，和测试值，，这里仅仅是用来定义的，方便后面的调用。
    sum = 0
    lens = len(y_predict)
    for i in range(lens):
        e = (y_predict[i] - y_test[i]) ** 2
        # errs.append(e)
        sum += e
    return math.sqrt(sum) / lens


if __name__ == '__main__':
    data = getData.GetData()
    data.getData()

    tra_model = tradaboost_r2.Tradaboost_r2()
    tra_model.fit(data)

    cdm_model = cdm.CDM()
    cdm_model.fit(data)

    fine_tune_model = fine_tune.Fine_tune()
    fine_tune_model.fit(data)

    x_test = data.target_x_test
    y_test = data.target_y_test

    y_pre1 = tra_model.predict(x_test)
    y_pre2 = cdm_model.predict(x_test)
    y_pre3 = fine_tune_model.predict(x_test)

    x_name = data.target_x_mean
    y_name = data.target_y_mean
    x_source = data.source_x
    y_source = data.source_y
    x_target = data.target_x
    y_target = data.target_y

    file1 = ["result_fig/cl对比.svg", "result_fig/cd对比.svg"]
    file2 = ["result_fig/cl预测结果.svg", "result_fig/cd预测结果.svg"]
    for i in range(y_pre1.shape[1]):
        plt.figure()
        plt.title(f"{y_name[i]}源域与目标域数据展示", fontsize=15)
        plt.xlabel(x_name[0], fontsize=15)
        plt.ylabel(y_name[i], fontsize=15)
        plt.plot(x_target, y_target[:, i], label="风洞数据", marker='.', color="black")
        plt.plot(x_source, y_source[:, i], label="XFLR数据", marker='.', color="blue")
        plt.legend(loc=0, frameon=False)
        plt.savefig(file1[i], bbox_inches='tight')
        plt.show()

        plt.figure()
        plt.title(f"{y_name[i]}预测对比", fontsize=15)
        plt.plot(x_test, y_pre1[:, i], label="TrAdaBoost.R2", linestyle='--', marker='s', markerfacecolor='white')
        plt.plot(x_test, y_pre2[:, i], label="CDM", linestyle='-.', marker='>', markerfacecolor='white')
        plt.plot(x_test, y_pre3[:, i], label="Fine_Tune", linestyle=':', marker='*', markerfacecolor='white')
        plt.plot(x_test, y_test[:, i], label="测试点", marker='.', color="black")
        # plt.plot(x_source, y_source[:, i], label="源域数据", marker='.')
        plt.xlabel(x_name[0], fontsize=15)
        plt.ylabel(y_name[i], fontsize=15)
        plt.legend(loc=0, frameon=False)
        plt.savefig(file2[i], bbox_inches='tight')
        plt.show()

    for i in range(len(y_name)):
        print('TrAdaBoost.R2方法的', y_name[i], "测试集的MSE误差为：", error(y_test[:, i], y_pre1[:, i]))
        print('TrAdaBoost.R2方法的', y_name[i], "测试集的R2误差为：", r2_score(y_test[:, i], y_pre1[:, i]))
        print('CDM.R2方法的', y_name[i], "测试集的MSE误差为：", error(y_test[:, i], y_pre2[:, i]))
        print('CDM.R2方法的', y_name[i], "测试集的R2误差为：", r2_score(y_test[:, i], y_pre2[:, i]))
        print('Fine_Tune方法的', y_name[i], "测试集的MSE误差为：", error(y_test[:, i], y_pre3[:, i]))
        print('Fine_Tune方法的', y_name[i], "测试集的R2误差为：", r2_score(y_test[:, i], y_pre3[:, i]))
        print("==" * 30)
